'''
Function:
    Implementation of Context Encoding Module
Author:
    Zhenchao Jin
'''
import copy
# import torch
# import torch.nn as nn
# import torch.nn.functional as F
import luojianet
import luojianet.nn as nn
import luojianet.ops as ops
from luojianet import Parameter, Tensor
from .encoding import Encoding
from ...backbones import BuildActivation, BuildNormalization, constructnormcfg


'''ContextEncoding'''
class ContextEncoding(nn.Module):
    def __init__(self, in_channels, num_codes, norm_cfg=None, act_cfg=None):
        super(ContextEncoding, self).__init__()
        # self.encoding_project = nn.Sequential(
        #     nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0, bias=False),
        #     BuildNormalization(constructnormcfg(placeholder=in_channels, norm_cfg=norm_cfg)),
        #     BuildActivation(act_cfg),
        # )
        self.encoding_project = nn.SequentialCell(
            nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0, pad_mode='pad', has_bias=False),
            BuildNormalization(constructnormcfg(placeholder=in_channels, norm_cfg=norm_cfg)),
            BuildActivation(act_cfg),
        )
        encoding_norm_cfg = copy.deepcopy(norm_cfg)
        encoding_norm_cfg['type'] = encoding_norm_cfg['type'].replace('2d', '1d')
        # self.encoding = nn.Sequential(
        #     Encoding(channels=in_channels, num_codes=num_codes),
        #     BuildNormalization(constructnormcfg(placeholder=num_codes, norm_cfg=encoding_norm_cfg)),
        #     BuildActivation(act_cfg),
        # )
        # self.fc = nn.Sequential(
        #     nn.Linear(in_channels, in_channels),
        #     nn.Sigmoid()
        # )
        self.encoding = nn.SequentialCell(
            # Encoding(channels=in_channels, num_codes=num_codes),
            Encoding(D=in_channels, K=num_codes),
            BuildNormalization(constructnormcfg(placeholder=num_codes, norm_cfg=encoding_norm_cfg)),
            BuildActivation(act_cfg),
        )
        self.fc = nn.SequentialCell(
            # nn.Linear(in_channels, in_channels),
            nn.Dense(in_channels, in_channels),
            nn.Sigmoid()
        )
    '''forward'''
    def forward(self, x):
        encoding_projection = self.encoding_project(x)
        encoding_feat = self.encoding(encoding_projection).mean(axis=1)
        # batch_size, channels, _, _ = x.size()
        batch_size, channels, _, _ = x.shape
        gamma = self.fc(encoding_feat)
        y = gamma.view(batch_size, channels, 1, 1)
        # output = F.relu_(x + x * y)
        output = ops.relu((x + x * y))
        return encoding_feat, output